import gradio as gr
from PIL import Image

from iflytech_assistant.ugc.avatar import base64_to_image
from iflytech_assistant.ugc.creator import generate_role_single_turn
from iflytech_assistant.ugc.role import Example, RoleMeta
from iflytech_assistant.ugc.role_play import generate


def create_role(
    state: dict,
    name: str,
    introduction: str,
    do_web_search: bool,
):
    role = generate_role_single_turn(
        name, do_web_search=do_web_search, generate_avatar=True
    )
    state[name] = role
    return (
        state,
        gr.Dropdown(label="选择角色", choices=state.keys()),
        role.introduction,
        role.category,
        role.appearance,
        role.profile,
        str(role.search_result),
        role.search_time + role.prefilling_time,
        role.generation_time,
        base64_to_image(role.avatar),
    )


def rewrite_message(
    state: dict,
    name: str,
    user_input: str,
    lines: str,
):
    role = state.get(name)
    role.lines = lines.split("\n")
    example: Example = generate(role, user_input=user_input, _print=False)
    return example.role_output, example.system_prompt + "\n" + example.user_prompt


def on_selected_role_change(state: dict, name: str):
    role: RoleMeta = state[name]
    return (
        role.name,
        role.introduction,
        role.category,
        role.appearance,
        str(role.search_result),
        base64_to_image(role.avatar),
    )


# Gradio 界面
with gr.Blocks() as demo:
    agent_state = gr.State(value={})

    with gr.Row():
        with gr.Column():
            gr.Markdown("### 角色配置")
            name = gr.Textbox(label="名字")
            do_web_search = gr.Checkbox(
                label="是否进行网络搜索", value=True, visible=False
            )
            profile = gr.Textbox(label="一句话描述")
            appearance = gr.Textbox(label="外貌特征")
            avatar = gr.Image(label="头像")
            introduction = gr.Textbox(label="角色介绍")
            category = gr.Textbox(label="角色类别", interactive=False)
            search_result = gr.Textbox(label="搜索结果", interactive=False, lines=3)
            create_button = gr.Button("创建角色")
            search_time = gr.Textbox(label="搜索时间", interactive=False)
            create_time = gr.Textbox(label="创建时间", interactive=False)

        with gr.Column():
            gr.Markdown("### 角色扮演")
            role_name = gr.Dropdown(label="选择角色", choices=[])
            lines = gr.Textbox(label="角色经典台词", lines=5)
            user_input = gr.Textbox(label="输入")
            output = gr.Textbox(label="输出", interactive=False, lines=5)
            rewrite_button = gr.Button("生成")
            raw_prompt = gr.Textbox(label="提示词", interactive=False, lines=10)

    create_button.click(
        fn=create_role,
        inputs=[agent_state, name, introduction, do_web_search],
        outputs=[
            agent_state,
            role_name,
            introduction,
            category,
            appearance,
            profile,
            search_result,
            search_time,
            create_time,
            avatar,
        ],
    )

    # on role_name change
    role_name.select(
        fn=on_selected_role_change,
        inputs=[agent_state, role_name],
        outputs=[name, introduction, category, appearance, search_result, avatar],
    )

    rewrite_button.click(
        fn=rewrite_message,
        inputs=[agent_state, role_name, user_input, lines],
        outputs=[output, raw_prompt],
    )


demo.launch(share=True)
